import torch.nn as nn
from torch.nn.utils import rnn

import chatbot.config as config


class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(num_embeddings=len(config.chatbot_ws_input),
                                      embedding_dim=config.chatbot_embedding_dim,
                                      padding_idx=config.chatbot_ws_input.PAD)

        self.gru = nn.GRU(input_size=config.chatbot_embedding_dim, batch_first=True,
                          num_layers=config.chatbot_encoder_num_layers,
                          hidden_size=config.chatbot_encoder_hidden_size)

    def forward(self, input, input_length):
        embedded = self.embedding(input)
        embedded = rnn.pack_padded_sequence(embedded, lengths=input_length.to('cpu'), batch_first=True)  # 打包
        output, hidden = self.gru(embedded)

        out, out_length = rnn.pad_packed_sequence(output, batch_first=True,
                                                  padding_value=config.chatbot_ws_input.PAD)  # 解包
        return out, hidden
